import os
import shutil
from patient_information import find_patient_files,load_patient_data,get_grade,get_murmur
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from python_speech_features import logfbank
from sklearn.model_selection import StratifiedKFold
from tqdm import tqdm
import wave
import librosa
import soundfile
from scipy import signal


def cut_copy_files_normal(data_directory: str, out_directory: str):
    for file_name in sorted(os.listdir(data_directory)):
        root, extension = os.path.splitext(file_name)
        patient_ID = root.split('_')[0]  # 获取病人ID
        grade = 'Absent'
        location = root.split('_')[5]
        recording, fs = librosa.load(os.path.join(data_directory, file_name), sr=4000)
        num_cut = len(recording) / (3 * 4000)  # 每个记录的片段数量
        start = 0
        end = start + 3 * fs
        cut = list()
        for num in range(int(num_cut)):  # 将每个片段写入对应的听诊区文件夹
            small = recording[start:end]
            cut.append(small)
            soundfile.write(
                out_directory + '/' + patient_ID + '_' + str(location) + '_' + str(grade) + '_' + str(num) + '.wav',
                cut[num], fs)
            start += 3 * fs
            end = start + 3 * fs





def cut_copy_files(data_directory: str, ident: str, out_directory: str) -> None:
    files = os.listdir(data_directory)
    for f in files:
        root, extension = os.path.splitext(f)
        if f.startswith(ident) :
            if extension == '.txt':
                _ = shutil.copy(os.path.join(data_directory, f), out_directory)
            elif extension == '.wav':
                #获取当前wav文件的ID 听诊区 等级
                with open(os.path.join(data_directory, ident+'.txt'), encoding="utf-8") as txt_f:

                    txt_data = txt_f.read()
                    patient_ID = root.split('_')[0]#获取病人ID
                    grade = get_clinical_grade(txt_data)
                    location = root.split('_')[5]
                recording, fs = librosa.load(os.path.join(data_directory, f), sr=4000)  # 分割（3s不重叠）

                num_cut = len(recording) / (3 * 4000)  # 每个记录的片段数量
                start = 0
                end = start+3*fs
                cut = list()
                for num in range(int(num_cut)):  # 将每个片段写入对应的听诊区文件夹
                    small = recording[start:end]
                    cut.append(small)
                    soundfile.write(out_directory + '/' + patient_ID + '_'+str(location)+'_' + str(grade) + '_' + str(num) + '.wav', cut[num], fs)
                    start += 3 * fs
                    end = start + 3 * fs
#只能处理确诊病例
def clinical_data(
    stratified_features: list,
    data_directory: str,#数据集路径
    out_directory: str,
    random_state: int,
):
    # Check if out_directory directory exists, otherwise create it.
    if not os.path.exists(out_directory):
        os.makedirs(out_directory)
    # else:
    #     exit()
        # shutil.rmtree(out_directory) #删除输出路径下的文件夹
        # os.makedirs(out_directory)
    # Get metadata
    patient_files = find_patient_files(data_directory)#获取排序后的病人text文件
    num_patient_files = len(patient_files)
    murmur_classes = ["Absent", "Soft", "Loud"]
    num_murmur_classes = len(murmur_classes)
    murmurs = list()
    patient_ID = list()
    for i in tqdm(range(num_patient_files)): #可视化遍历进度条，遍历每个病人的txt文件
        # Load the current patient data and recordings.
        _, name = os.path.split(patient_files[i])
        txt_name, extension = os.path.splitext(name)
        fn = open(patient_files[i], encoding="utf-8")  # 打开文件
        current_patient_data = fn.read()  # 读出整个文件

        # current_patient_data = load_patient_data(patient_files[i]) #获取text文本内容
        murmur_unkonwn=get_clinical_murmur(current_patient_data)
        #跳过没有杂音等级的数据
        if murmur_unkonwn == 'nan':
            continue
        current_ID = txt_name.split("_")[0] #获取ID
        patient_ID.append(current_ID)#添加ID
        current_murmur = np.zeros(num_murmur_classes, dtype=int)
        murmur = get_clinical_grade(current_patient_data)  # 获取杂音标签
        if murmur in murmur_classes:
            j = murmur_classes.index(murmur)
            current_murmur[j] = 1            #"Aresent"=100,soft=010
        murmurs.append(current_murmur)   #保存杂音标签
    patient_ID = np.vstack(patient_ID)
    num_murmurs=len(murmurs);#筛除unknown后的患者数量
    print(len(murmurs))
    murmurs = np.vstack(murmurs);

    #指标
    patients_pd = pd.DataFrame(patient_ID,columns=['ID'],)
    murmurs_pd = pd.DataFrame(murmurs, columns=murmur_classes)
    complete_pd = pd.concat([patients_pd, murmurs_pd], axis=1)  # 拼接表格（ID，杂音label）
    print(complete_pd)
    # complete_pd["ID"] = complete_pd["ID"].astype(int).astype(str)
    # Split data
    complete_pd["stratify_column"] = (
        complete_pd[stratified_features].astype(str).agg("-".join, axis=1)# 添加一列：Absent-Soft-Loud
    )
    print(complete_pd)
    if not os.path.exists(os.path.join(out_directory, "wav_cut_data")):
        os.makedirs(os.path.join(out_directory, "wav_cut_data"))
    for f in complete_pd["ID"]:
        print(f)
        cut_copy_files(
            data_directory,
            f,
            os.path.join(out_directory, "wav_cut_data/"),
        )


def FrameTimeC(frameNum, frameLen, inc, fs):
    ll = np.array([i for i in range(frameNum)])
    return ((ll - 1) * inc + frameLen / 2) / fs
def FrequencyScale(nfilt,fs):
    high_freq_mel = (2595 * np.log10(1 + (fs / 2) / 700))  # 求最高hz频率对应的mel频率
# 我们要做40个滤波器组，为此需要42个点，这意味着在们需要low_freq_mel和high_freq_mel之间线性间隔40个点
    mel_points = np.linspace(0, high_freq_mel, nfilt + 2)  # 在mel频率上均分成42个点
    hz_points = (700 * (10 ** (mel_points / 2595) - 1))  # 将mel频率再转到hz频率
    return hz_points

def Log_Mel_Spec(data_directory ):
    logmelspec = list()
    for f in sorted(os.listdir(data_directory)):
        root, extension = os.path.splitext(f)
        if extension == '.wav':
            x, fs = librosa.load(os.path.join(data_directory, f), sr=4000)
            b, a = signal.butter(4, [0.01, 0.375], 'bandpass')  # 配置滤波器 8 表示滤波器的阶数
            x = signal.filtfilt(b, a, x)  # data为要过滤的信号
            x = x - np.mean(x)
            x = x / np.max(np.abs(x))
            fbank_feat = logfbank(x,fs,winlen=0.025,winstep=0.0125,nfilt=32,nfft=512,lowfreq=0,highfreq=800)
            fbank_feat=fbank_feat.T
            logmelspec.append(fbank_feat)
        else:
            continue
    return np.array(logmelspec)
# Get murmur from patient data.
def get_clinical_murmur(data):
    murmur = None
    for text in data.split("\n"):
        if text.startswith("杂音等级"):
            murmur = text.split("：")[1]
    if murmur is None:
        raise ValueError(
            "No murmur available. Is your code trying to load labels from the hidden data?"
        )
    return murmur

def get_clinical_grade(data):
    grade = None
    for text in data.split("\n"):
        if text.startswith("杂音等级"):
            grading = text.split("：")[1].strip()
            if grading == 'I' or grading == 'II':
                grade = 'Soft'
            elif grading == 'Absent':
                grade = 'Absent'
            else:
                grade = 'Loud'
    return grade

def get_label(data_directory ):
    label=list()
    location=list()
    id=list()

    for f in sorted(os.listdir(data_directory)):
        root, extension = os.path.splitext(f)
        if extension == '.wav':
            the_label = root.split("_")[2].strip()
            if the_label == 'Absent':
                grade = 0
            elif the_label == 'Soft':
                grade = 1
            elif the_label == 'Loud':
                grade = 2
            label.append(grade)

            the_location=root.split("_")[1].strip()
            location.append(the_location)

            the_id=root.split("_")[0].strip()
            id.append(the_id)


    return np.array(label),np.array(location),np.array(id)

def get_index(data_directory ):
    index = list()
    i=0
    for f in sorted(os.listdir(data_directory)):
        root, extension = os.path.splitext(f)
        if extension == '.wav':
            index.append(i)
            i=i+1
        else:
            continue
    return np.array(index)


if __name__ == "__main__":
    data_directory = "E:/HZH/heart_data/clinical_heart_data"
    normal_data_directory ="E:/HZH/heart_data/clinical_heart_data_normal"
    normal_out_directory = "clinic_data/wav_cut_data/"
    out_directory="clinic_data"
    random_state=5678
    stratified_features = ["Absent", "Soft", "Loud"]
    #数据切分
    # clinical_data(stratified_features,data_directory,out_directory,random_state)
    # cut_copy_files_normal(normal_data_directory,normal_out_directory)


    test_data_directory =  os.path.join(out_directory, "wav_cut_data")
    print(test_data_directory)
    #特征 标签提取
    label_directory = os.path.join(out_directory, "label")
    logmel_directory = os.path.join(out_directory, "logmel")
    if not os.path.exists(logmel_directory):
        os.makedirs(logmel_directory)
        test_feature = Log_Mel_Spec(test_data_directory)
        np.save(logmel_directory + r'/test_feature.npy', test_feature)
    if not os.path.exists(label_directory):
        os.makedirs(label_directory)
        test_label, test_location, test_id = get_label(test_data_directory)
        np.save(label_directory + r'/test_label.npy', test_label)
        np.save(label_directory + r'/test_location.npy', test_location)
        np.save(label_directory + r'/test_id.npy', test_id)
        test_index = get_index(test_data_directory)
        np.save(label_directory + r'/test_index.npy', test_index)










